use std::env;

use anyhow::Result;
use clap::{Arg, ArgAction, Command};
use rmcp::{
    ServerHandler, ServiceExt,
    handler::server::{tool::ToolRouter, wrapper::Parameters},
    model::{Implementation, ProtocolVersion, ServerCapabilities, ServerInfo},
    schemars, tool, tool_handler, tool_router,
    transport::stdio,
};

mod tool;

fn main() {
    // 环境变量
    let matches = Command::new("rig-search")
        .about("基于LLM解释代码，并向量化搜索的工具")
        .arg(
            Arg::new("project-path")
                .long("project-path")
                .value_parser(clap::value_parser!(String))
                .help("项目路径"),
        )
        .arg(
            Arg::new("class-name")
                .long("class-name")
                .value_parser(clap::value_parser!(String))
                .help("类的全名"),
        )
        .arg(
            Arg::new("method-name")
                .long("method-name")
                .value_parser(clap::value_parser!(String))
                .help("方法名"),
        )
        .arg(
            Arg::new("interface-name")
                .long("interface-name")
                .value_parser(clap::value_parser!(String))
                .help("接口名"),
        )
        .arg(
            Arg::new("mcp")
                .long("mcp")
                .num_args(0)
                .action(ArgAction::SetTrue)
                .help("启动mcp，或者设置JAVA_PARSER_MCP_ALLOW=true"),
        )
        .arg(
            Arg::new("relate-code")
                .long("relate-code")
                .num_args(0)
                .action(ArgAction::SetTrue)
                .help("输出相关的代码"),
        )
        .arg(
            Arg::new("spring-url")
                .long("spring-url")
                .num_args(0)
                .action(ArgAction::SetTrue)
                .help("输出Spring项目中所有url和实现代码的reference"),
        )
        .arg(Arg::new("filename").value_parser(clap::value_parser!(String)))
        .get_matches();

    // 检查是否是mcp
    let mcp_env_allow = env::var("JAVA_PARSER_MCP_ALLOW")
        .unwrap_or("false".to_string())
        .parse::<bool>()
        .unwrap_or(false);
    if matches.get_flag("mcp") || mcp_env_allow {
        let rt = tokio::runtime::Runtime::new().unwrap();
        let result = rt.block_on(start_mcp());
        return;
    }

    let default_path = ".".to_string().to_owned();
    let project_path: &String = matches
        .get_one::<String>("project-path")
        .unwrap_or(&default_path);

    let method_name = matches.get_one::<String>("method-name");

    // 查询相关代码
    if matches.get_flag("relate-code") {
        if let Some(class_name) = matches.get_one::<String>("class-name") {
            let relate_code = tool::relateCodeByClassAndMethod::recursion_get_used_code(
                project_path.as_str(),
                class_name,
                method_name.map(|item| item.as_str()),
                false,
            );
            println!("{relate_code}");
            return;
        } else {
            panic!("需要设置类名");
        }
    }

    // 查询类
    if let Some(class_name) = matches.get_one::<String>("class-name") {
        let results = tool::byClassAndMethod::by_class_name_and_method(
            project_path.as_str(),
            class_name,
            method_name,
        );
        println!(
            "{}",
            serde_json::to_string(&results).unwrap_or("[]".to_string())
        );
        return;
    }

    // 查询接口
    if let Some(interface_name) = matches.get_one::<String>("interface-name") {
        let results =
            tool::byInterfaceName::by_interface_name(project_path.as_str(), interface_name);
        println!(
            "{}",
            serde_json::to_string(&results).unwrap_or("[]".to_string())
        );
        return;
    }

    // 获取spring项目中所有url
    if matches.get_flag("spring-url") {
        let result: Vec<(String, String, String)> =
            tool::base::get_url_and_class_from_project(project_path.as_str());

        for item in result {
            println!("{} {} {}", item.0, item.1, item.2)
        }
        return;
    }
}

// MCP
#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct ClassRequest {
    #[schemars(description = "项目路径")]
    pub project_path: String,

    #[schemars(description = "类的全路径名")]
    pub class_name: String,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct ClassAndMethodRequest {
    #[schemars(description = "项目路径")]
    pub project_path: String,

    #[schemars(description = "类的全路径名")]
    pub class_name: String,

    #[schemars(description = "方法名")]
    pub method_name: String,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct RelateCodeByClassAndMethodRequest {
    #[schemars(description = "项目路径")]
    pub project_path: String,

    #[schemars(description = "类的全路径名")]
    pub class_name: String,

    #[schemars(description = "方法名")]
    pub method_name: String,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct InterfaceRequest {
    #[schemars(description = "项目路径")]
    pub project_path: String,

    #[schemars(description = "查找接口的直接实现类，接口名需要全路径")]
    pub interface_name: String,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct SpringUrlRequest {
    #[schemars(description = "项目路径")]
    pub project_path: String,
}

#[derive(Debug, Clone)]
pub struct ProjectParser {
    tool_router: ToolRouter<Self>,
}

#[tool_router]
impl ProjectParser {
    #[allow(dead_code)]
    pub fn new() -> Self {
        Self {
            tool_router: Self::tool_router(),
        }
    }

    #[tool(description = "根据类/接口的全名查询相关代码")]
    fn by_class(
        &self,
        Parameters(ClassRequest {
            project_path,
            class_name,
        }): Parameters<ClassRequest>,
    ) -> String {
        let results = tool::byClassAndMethod::by_class_name_and_method(
            project_path.as_str(),
            class_name.as_str(),
            None,
        );
        serde_json::to_string(&results).unwrap_or("[]".to_string())
    }

    #[tool(description = "根据类的全名和方法名查询方法体")]
    fn by_class_and_method(
        &self,
        Parameters(ClassAndMethodRequest {
            project_path,
            class_name,
            method_name,
        }): Parameters<ClassAndMethodRequest>,
    ) -> String {
        let results = tool::byClassAndMethod::by_class_name_and_method(
            project_path.as_str(),
            class_name.as_str(),
            Some(&method_name),
        );
        serde_json::to_string(&results).unwrap_or("[]".to_string())
    }

    #[tool(description = "根据类的全名和方法名查询所有相关的代码")]
    fn relate_code_by_class_and_method(
        &self,
        Parameters(RelateCodeByClassAndMethodRequest {
            project_path,
            class_name,
            method_name,
        }): Parameters<RelateCodeByClassAndMethodRequest>,
    ) -> String {
        let relate_code = tool::relateCodeByClassAndMethod::recursion_get_used_code(
            project_path.as_str(),
            class_name.as_str(),
            Some(method_name.as_str()),
            false,
        );
        serde_json::to_string(&relate_code).unwrap_or("[]".to_string())
    }

    #[tool(description = "根据接口的全名查询实现的类")]
    fn by_interface_name(
        &self,
        Parameters(InterfaceRequest {
            project_path,
            interface_name,
        }): Parameters<InterfaceRequest>,
    ) -> String {
        let implements_class_name: Vec<String> =
            tool::byInterfaceName::by_interface_name(project_path.as_str(), &interface_name);
        implements_class_name.join("\n")
    }

    #[tool(description = "查询当前Spring项目的所有接口。按http方法，url，类加方法的格式返回")]
    fn all_url(
        &self,
        Parameters(SpringUrlRequest { project_path }): Parameters<SpringUrlRequest>,
    ) -> String {
        let result: Vec<(String, String, String)> =
            tool::base::get_url_and_class_from_project(project_path.as_str());
        let lines: Vec<String> = result
            .iter()
            .map(|item| format!("{} {} {}", item.0, item.1, item.2))
            .collect();
        lines.join("\n")
    }
}

#[tool_handler]
impl ServerHandler for ProjectParser {
    fn get_info(&self) -> ServerInfo {
        ServerInfo {
            protocol_version: ProtocolVersion::V_2024_11_05,
            instructions: Some("Java 工程的工具类,通过类名和方法名精确查询相关代码".into()),
            capabilities: ServerCapabilities::builder().enable_tools().build(),
            server_info: Implementation::from_build_env(),
        }
    }
}

async fn start_mcp() -> Result<()> {
    let service = ProjectParser::new()
        .serve(stdio())
        .await
        .inspect_err(|e| eprintln!("{e}"))?;

    service.waiting().await?;
    Ok(())
}
